# coding: utf-8 
"""
@Time    : 2024/8/14 16:51
@Author  : Y.H LEE
"""
import torch
from torch import nn
import torch.nn.functional as F

from sys_params import device
from utils.metrics import *

"""
MS Task
"""


class LSTM(nn.Module):
    def __init__(self, in_dim, hid_dim, num_layers, out_dim):
        super().__init__()
        self.hid_dim = hid_dim
        self.num_layers = num_layers

        self.lstm = nn.LSTM(in_dim, hid_dim, num_layers, batch_first=True)
        self.linear = nn.Linear(hid_dim, out_dim)

    def forward(self, data):
        x, y = data[0].unsqueeze(-1), data[1]

        h_0 = torch.zeros(self.num_layers, x.shape[0], self.hid_dim).to(device)
        c_0 = torch.zeros(self.num_layers, x.shape[0], self.hid_dim).to(device)

        output, (_, _) = self.lstm(x, (h_0, c_0))
        batch_size, timestamp, hidden_dim = output.shape
        output = output[:, -1, :].squeeze()

        output = self.linear(output)

        mse_loss, mape_loss, sgcc = self.get_metric(output.squeeze(), y)

        return mse_loss, mape_loss, sgcc

    def get_metric(self, output, y):
        # print(output)
        mse_loss = MSE(output, y)
        mape_loss = MAPE(output, y)
        sgcc = SGCC(output, y)

        return mse_loss, mape_loss, sgcc
